# main_trajectory.py
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os

# 导入自定义的场地模块
from badminton_court_3d import BadmintonCourt3D

def read_trajectory_from_txt(file_path, delimiter=','):
    """
    从.txt文件读取羽毛球轨迹数据
    
    参数:
    file_path: .txt文件路径
    delimiter: 分隔符，默认为逗号
    """
    try:
        # 读取.txt文件
        df = pd.read_csv(file_path, delimiter=delimiter)
        print(f"成功读取文件: {file_path}")
        print(f"数据形状: {df.shape}")
        print(f"列名: {list(df.columns)}")
        return df
    except Exception as e:
        print(f"读取文件时出错: {e}")
        return None

def read_trajectory_auto_detect(file_path):
    """
    自动检测文件格式并读取数据
    """
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        return None
    
    # 根据文件扩展名选择读取方式
    if file_path.endswith('.txt'):
        # 尝试不同的分隔符
        for delimiter in [',', '\t', ';', ' ']:
            try:
                df = pd.read_csv(file_path, delimiter=delimiter)
                if len(df.columns) >= 4:  # 至少需要4列数据
                    print(f"使用分隔符 '{delimiter}' 成功读取文件")
                    return df
            except:
                continue
    elif file_path.endswith('.csv'):
        try:
            df = pd.read_csv(file_path)
            return df
        except Exception as e:
            print(f"读取CSV文件时出错: {e}")
            return None
    
    print("无法读取文件，请检查文件格式")
    return None

def transform_coordinates(df, x_scale=1.0, y_scale=1.0, z_invert=False):
    """
    对坐标进行变换：x,y值乘以缩放因子，z值取反
    
    参数:
    df: 原始数据DataFrame
    x_scale: x坐标缩放因子
    y_scale: y坐标缩放因子
    z_invert: 是否对z坐标取反
    """
    df_transformed = df.copy()
    
    # 应用坐标变换
    if 'x' in df.columns:
        df_transformed['x'] = df_transformed['x'] * x_scale
        print(f"x坐标缩放: 乘以 {x_scale}")
    
    if 'y' in df.columns:
        df_transformed['y'] = df_transformed['y'] * y_scale
        print(f"y坐标缩放: 乘以 {y_scale}")
    
    if 'z' in df.columns and z_invert:
        df_transformed['z'] = -df_transformed['z']
        print("z坐标取反")
    
    return df_transformed

def plot_trajectory_with_court(txt_file_path, transform_coords=True, 
                              x_scale=1.0, y_scale=1.0, z_invert=False, 
                              **court_kwargs):
    """
    主函数：从.txt文件读取数据并绘制带轨迹的羽毛球场地
    
    参数:
    txt_file_path: .txt文件路径
    transform_coords: 是否进行坐标变换
    x_scale: x坐标缩放因子
    y_scale: y坐标缩放因子
    z_invert: 是否对z坐标取反
    court_kwargs: 传递给场地创建的参数
    """
    # 读取轨迹数据
    df = read_trajectory_auto_detect(txt_file_path)
    
    if df is None or df.empty:
        print("无法读取数据或数据为空")
        return
    
    # 显示原始数据前几行
    print("\n=== 原始数据前5行 ===")
    print(df.head())
    print(f"\n原始数据列名: {list(df.columns)}")
    
    # 坐标变换
    if transform_coords:
        print("\n=== 应用坐标变换 ===")
        df_transformed = transform_coordinates(df, x_scale, y_scale, z_invert)
        
        # 显示变换后的数据
        print("\n=== 变换后数据前5行 ===")
        print(df_transformed.head())
    else:
        df_transformed = df.copy()
        print("跳过坐标变换")
    
    # 创建3D羽毛球场地
    court = BadmintonCourt3D()
    
    # 设置默认参数（可被court_kwargs覆盖）
    default_kwargs = {
        'figsize': (10, 18),
        'view_elev': 20,
        'view_azim': -45,
        'bg_color': 'black',
        'court_color': 'green'
    }
    default_kwargs.update(court_kwargs)
    
    fig, ax = court.create_court(**default_kwargs)
    
    # 添加羽毛球轨迹（使用变换后的坐标）
    scatter = court.add_shuttlecock_trajectory(ax, df_transformed)
    
    # 如果有轨迹数据，添加颜色条
    if scatter is not None:
        cbar = plt.colorbar(scatter, ax=ax, pad=0.1)
        cbar.set_label('Frame Index', rotation=270, labelpad=15)
    
    # 添加图例
    ax.legend(loc='upper left', bbox_to_anchor=(0, 1))
    
    # 更新标题以包含变换信息
    title = f"Badminton Court with Transformed Trajectory\n{os.path.basename(txt_file_path)}"
    if transform_coords:
        title += f"\nTransform: x×{x_scale}, y×{y_scale}, z×{-1 if z_invert else 1}"
    
    plt.title(title, fontsize=14, color='red', pad=20)
    
    # 显示图形
    plt.tight_layout()
    plt.show()
    
    # 打印统计信息（变换后的）
    print_trajectory_statistics(df_transformed, "变换后")

def print_trajectory_statistics(df, label="数据"):
    """打印轨迹统计信息"""
    if 'visible' in df.columns:
        df_visible = df[df['visible'] == 1]
    else:
        df_visible = df
    
    print(f"\n=== {label}轨迹统计信息 ===")
    print(f"总数据点: {len(df)}")
    print(f"可见数据点: {len(df_visible)}")
    
    if len(df_visible) > 0:
        print(f"起点坐标: ({df_visible.iloc[0]['x']:.2f}, {df_visible.iloc[0]['y']:.2f}, {df_visible.iloc[0]['z']:.2f})")
        print(f"终点坐标: ({df_visible.iloc[-1]['x']:.2f}, {df_visible.iloc[-1]['y']:.2f}, {df_visible.iloc[-1]['z']:.2f})")
        print(f"最大高度: {df_visible['z'].max():.2f} m")
        print(f"最小高度: {df_visible['z'].min():.2f} m")
        print(f"X范围: [{df_visible['x'].min():.2f}, {df_visible['x'].max():.2f}]")
        print(f"Y范围: [{df_visible['y'].min():.2f}, {df_visible['y'].max():.2f}]")
        print(f"Z范围: [{df_visible['z'].min():.2f}, {df_visible['z'].max():.2f}]")

def compare_original_vs_transformed(txt_file_path, x_scale=1.0, y_scale=1.0, z_invert=False):
    """
    对比显示原始坐标和变换后坐标的轨迹
    """
    # 读取数据
    df = read_trajectory_auto_detect(txt_file_path)
    if df is None:
        return
    
    # 应用变换
    df_transformed = transform_coordinates(df, x_scale, y_scale, z_invert)
    
    # 创建对比图形
    fig = plt.figure(figsize=(20, 8))
    
    # 原始坐标轨迹
    ax1 = fig.add_subplot(121, projection='3d')
    court = BadmintonCourt3D()
    court._setup_axes(ax1, 'black')  # 只设置坐标轴，不绘制完整场地
    scatter1 = court.add_shuttlecock_trajectory(ax1, df)
    ax1.set_title('原始坐标轨迹', fontsize=14, color='white', pad=20)
    
    # 变换后坐标轨迹
    ax2 = fig.add_subplot(122, projection='3d')
    court._setup_axes(ax2, 'black')
    scatter2 = court.add_shuttlecock_trajectory(ax2, df_transformed)
    ax2.set_title(f'变换后坐标轨迹\nx×{x_scale}, y×{y_scale}, z×{-1 if z_invert else 1}', 
                  fontsize=14, color='white', pad=20)
    
    plt.tight_layout()
    plt.show()
    
    # 打印对比统计信息
    print_trajectory_statistics(df, "原始")
    print_trajectory_statistics(df_transformed, "变换后")

def batch_process_with_transform(folder_path, x_scale=2.0, y_scale=2.0, z_invert=True, file_pattern="*.txt"):
    """
    批量处理文件夹中的多个轨迹文件并应用坐标变换
    """
    import glob
    
    txt_files = glob.glob(os.path.join(folder_path, file_pattern))
    
    if not txt_files:
        print(f"在文件夹 {folder_path} 中未找到 {file_pattern} 文件")
        return
    
    print(f"找到 {len(txt_files)} 个轨迹文件，应用坐标变换:")
    print(f"  x缩放: {x_scale}, y缩放: {y_scale}, z取反: {z_invert}")
    
    for i, file_path in enumerate(txt_files):
        print(f"\n处理文件 {i+1}/{len(txt_files)}: {os.path.basename(file_path)}")
        plot_trajectory_with_court(file_path, True, x_scale, y_scale, z_invert)

# 示例使用
if __name__ == "__main__":
    # 替换为您的.txt文件路径
    txt_file_path = "../../3D坐标/badminton_6v/case_0_3d.txt"
    
    if os.path.exists(txt_file_path):
        # 方法1: 应用坐标变换并显示
        print("=== 方法1: 应用坐标变换 ===")
        plot_trajectory_with_court(
            txt_file_path, 
            transform_coords=False,
            x_scale=1.0, 
            y_scale=1.0, 
            z_invert=False
        )
        
        # 方法2: 对比显示原始和变换后的轨迹
        print("\n=== 方法2: 对比显示 ===")
        compare_original_vs_transformed(txt_file_path)
        
        # 方法3: 不应用变换（原始坐标）
        print("\n=== 方法3: 原始坐标 ===")
        plot_trajectory_with_court(
            txt_file_path, 
            transform_coords=False
        )
        
    else:
        print(f"文件 {txt_file_path} 不存在")
        
        # 显示当前目录下的.txt文件
        current_dir = os.getcwd()
        txt_files = [f for f in os.listdir(current_dir) if f.endswith('.txt')]
        if txt_files:
            print(f"\n当前目录下的.txt文件:")
            for f in txt_files:
                print(f"  - {f}")
            # 使用第一个找到的文件
            plot_trajectory_with_court(txt_files[0], False, 1.0, 1.0, False)